from flask import Flask, request, jsonify
from flask_cors import CORS
import numpy as np
from tensorflow.keras.models import load_model
from PIL import Image
import io

# 创建Flask应用（避免使用默认名称app）
flask_app = Flask(__name__)
CORS(flask_app)

# 提前加载模型（避免在路由中重复加载）
try:
    model = load_model('mnist_model.h5')
except:
    print("模型加载失败，请检查mnist_model.h5是否存在")
    exit(1)

def preprocess_image(image):
    """统一的图像预处理函数"""
    image = image.convert('L').resize((28, 28))
    arr = np.array(image).astype('float32') / 255.0
    return np.expand_dims(arr, axis=(0, -1))  # 同时添加batch和channel维度

@flask_app.route('/predict', methods=['POST'])
def handle_predict():
    """处理预测请求"""
    if 'file' not in request.files:
        return jsonify({"error": "No file part"}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    try:
        img = Image.open(io.BytesIO(file.read()))
        processed = preprocess_image(img)
        pred = model.predict(processed)
        return jsonify({
            "prediction": int(np.argmax(pred)),
            "confidence": float(np.max(pred)),
            "status": "success"
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# 确保路由在直接运行时注册
if __name__ == '__main__':
    flask_app.run(host='0.0.0.0', port=5001, debug=False)  # 关闭debug避免自动重载